﻿#include <iostream>
#include <opencv2/opencv.hpp>

// 判断颜色是否为红色
bool isRedColor(const cv::Mat& hsvFrame, int x, int y) {
    cv::Vec3b color = hsvFrame.at<cv::Vec3b>(y, x);
    return color[0] < 10 && color[1] > 50 && color[2] > 50;
}

// 判断颜色是否为蓝色
bool isBlueColor(const cv::Mat& hsvFrame, int x, int y) {
    cv::Vec3b color = hsvFrame.at<cv::Vec3b>(y, x);
    return color[0] > 100 && color[1] > 50 && color[2] < 50;
}

// 灯条识别函数
void detectLightBars(cv::Mat frame, const std::string& color1) {
    cv::Mat hsvFrame;
    cv::cvtColor(frame, hsvFrame, cv::COLOR_BGR2HSV);

    // 二值化处理
    cv::Mat binaryFrame;
    cv::adaptiveThreshold(cv::Mat(hsvFrame.clone()), binaryFrame, 255, cv::ADAPTIVE_THRESH_GAUSSIAN_C, cv::THRESH_BINARY, 11, 2);

    // 创建结构元素用于膨胀和腐蚀
    cv::Mat structuringElement = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(15, 15));

    // 进行膨胀和腐蚀操作
    cv::dilate(binaryFrame, binaryFrame, structuringElement);
    cv::erode(binaryFrame, binaryFrame, structuringElement);
    cv::dilate(binaryFrame, binaryFrame, structuringElement);

    // 轮廓检测
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(binaryFrame, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    std::vector<cv::Rect> lightBars;

    for (const auto& contour : contours) {
        cv::Rect boundingRect = cv::boundingRect(contour);
        double area = cv::contourArea(contour);
        // 降低面积阈值以允许更小的灯条被检测
        if (area > 50) { // 面积阈值可以根据实际情况调整
            bool isColor = false;
            for (int y = boundingRect.y; y < boundingRect.y + boundingRect.height; y++) {
                for (int x = boundingRect.x; x < boundingRect.x + boundingRect.width; x++) {
                    if (x < frame.cols && y < frame.rows) {
                        cv::Vec3b color = hsvFrame.at<cv::Vec3b>(y, x);
                        if ((color1 == "red" && isRedColor(hsvFrame, x, y)) || (color1 == "blue" && isBlueColor(hsvFrame, x, y))) {
                            isColor = true;
                            break;
                        }
                    }
                }
                if (isColor) break;
            }
            if (isColor) {
                lightBars.push_back(boundingRect);
            }
        }
    }

    // 绘制灯条
    for (const auto& lightBar : lightBars) {
        cv::rectangle(frame, lightBar, cv::Scalar(0, 255, 0), 2);
    }

    cv::imshow("Light Bars", frame);
}

int main() {
    std::string videoPath = "C:\\Users\\aaa\\Pictures\\Screenshots\\步兵素材红车旋转-ev--3.MOV";
    cv::VideoCapture video(videoPath);

    if (!video.isOpened()) {
        std::cout << "Could not open the video." << std::endl;
        return -1;
    }

    cv::Mat frame;
    std::string enermycolor = "red";
    while (true) {
        video >> frame;
        if (frame.empty())
            break;

        detectLightBars(frame, enermycolor);

        if (cv::waitKey(30) >= 0)
            break;
    }

    video.release();
    cv::destroyAllWindows();

    return 0;
}